#define _CRT_SECURE_NO_WARNINGS

#include "sm2.h"

#include <string.h>
#include <malloc.h>
#include <stdlib.h>

#include "pointoperation.h"
#include "kdf.h"
#include "sm3.h"

/**
 * 初始化SM2椭圆曲线参数.
 *
 * \param ecc_param
 */
void sm2_param_init(param sm2_param)
{
    // 分配空间
    mpz_inits(sm2_param->p, sm2_param->a, sm2_param->b, sm2_param->n, NULL);

    // 初始化
    // 采用16进制字符串
    int base = 16;

    // p = 8542D69E 4C044F18 E8B92435 BF6FF7DE 45728391 5C45517D 722EDB8B 08F1DFC3
    mpz_set_str(sm2_param->p, "8542D69E4C044F18E8B92435BF6FF7DE457283915C45517D722EDB8B08F1DFC3", base);

    // a = 787968B4 FA32C3FD 2417842E 73BBFEFF 2F3C848B 6831D7E0 EC65228B 3937E498
    mpz_set_str(sm2_param->a, "787968B4FA32C3FD2417842E73BBFEFF2F3C848B6831D7E0EC65228B3937E498", base);

    // b = 63E4C6D3 B23B0C84 9CF84241 484BFE48 F61D59A5 B16BA06E 6E12D1DA 27C5249A
    mpz_set_str(sm2_param->b, "63E4C6D3B23B0C849CF84241484BFE48F61D59A5B16BA06E6E12D1DA27C5249A", base);

    // n = 8542D69E 4C044F18 E8B92435 BF6FF7DD 29772063 0485628D 5AE74EE7 C32E79B7
    mpz_set_str(sm2_param->n, "8542D69E4C044F18E8B92435BF6FF7DD297720630485628D5AE74EE7C32E79B7", base);
}

/**
 * 释放SM2椭圆曲线参数的空间.
 *
 * \param sm2_param
 */
void sm2_param_clear(param sm2_param)
{
    // 分配空间
    mpz_clears(sm2_param->p, sm2_param->a, sm2_param->b, sm2_param->n, NULL);
}

/**
 * 根据私钥private_key 和 基点G生成公钥.
 *
 * \param public_key
 * \param private_key
 * \param sm2_param
 * \param base_point
 */
void public_key_gen(point public_key, mpz_t private_key, param sm2_param, point base_point)
{
    point_mul(public_key, base_point, private_key, sm2_param->a, sm2_param->p);
}

/**
 * SM2加密.
 *
 * \param ct 密文
 * \param ct_len 密文长度
 * \param sm2_param SM2的参数
 * \param public_key 公钥
 * \param base_point 基点G
 * \param msg 消息（明文）
 * \param msg_len 明文（明文）的长度
 */
void sm2_encrypt(uint8_t *ct, size_t ct_len, param sm2_param, point public_key, point base_point, uint8_t *msg, size_t msg_len)
{
    // 初始化
    // 字符串使用16进制
    size_t base = 16;
    // 坐标的长度 48个16进制数
    // size_t position_len = 48;
    size_t position_len = 24;
    // mpz转换成字符串之后的长度，符号+结尾 （+2字节）。但是这里不考虑符号，所以只+1
    size_t mpz_str_len = position_len + 1;

    size_t ct_index = 0;

    // 随机数
    // 384F3035 3073AEEC E7A16543 30A96204 D37982A3 E15B2CB5
    mpz_t rand_k;
    mpz_init_set_str(rand_k, "384F30353073AEECE7A1654330A96204D37982A3E15B2CB5", base);

    // 1.生成密文C1
    // 1.1.初始化多倍点 k[G]
    point k_g;
    point_init(k_g);

    // 1.2.计算 (x1, y2) = k[G] =rand_k[G]
    point_mul(k_g, base_point, rand_k, sm2_param->a, sm2_param->p);

    // 1.3.计算 C1 = 04 || x1 || y1
    ct[ct_index] = 0x04;
    ct_index = ct_index + 1;

    // mpz_to_fixed_length_str(ct + 2, mpz_str_len, k_g->x, base);
    mpz_to_byte_array(ct + ct_index, position_len, k_g->x);
    ct_index = ct_index + position_len;

    // mpz_to_fixed_length_str(ct + 2 + position_len, mpz_str_len, k_g->y, base);
    mpz_to_byte_array(ct + ct_index, position_len, k_g->y);
    ct_index = ct_index + position_len;

    // print_chars_by_hex("打印ct1", ct + 1, 2 + position_len * 2);

    // 2.生成密文C2
    // 2.1.计算多倍点 (x2, y2) = k[Pb] = rand_k[Pb]
    point k_public_key;
    point_init(k_public_key);
    point_mul(k_public_key, public_key, rand_k, sm2_param->a, sm2_param->p);

    size_t x2_y2_len = 2 * position_len;
    uint8_t *x2_y2 = (uint8_t *)malloc(sizeof(uint8_t) * x2_y2_len);

    // mpz_to_fixed_length_str(x2_y2, mpz_str_len, k_public_key->x, base);
    mpz_to_byte_array(x2_y2, position_len, k_public_key->x);
    // mpz_to_fixed_length_str(x2_y2 + position_len, mpz_str_len, k_public_key->y, base);
    mpz_to_byte_array(x2_y2 + position_len, position_len, k_public_key->y);
    // print_chars_by_hex("x2_y2", x2_y2, x2_y2_len);

    // 2.2.使用密钥派生函数，生成 t = KDF(x2 || y2, len) ,其中 len = klen/8，klen是8的倍数
    uint8_t *ct2 = (uint8_t *)malloc(sizeof(uint8_t) * msg_len);

    kdf(ct2, msg_len, x2_y2, x2_y2_len);
    // print_chars_by_hex("r = KDF(x2 ‖ y2, klen)", ct2, msg_len);

    // 2.3.计算 C2 = M ^ t
    for (size_t i = 0; i < msg_len; i++)
    {
        ct2[i] = msg[i] ^ ct2[i];
    }
    // print_chars_by_hex("ct2", ct2, msg_len);

    // 3.生成密文C3
    // C3 = Hash(x2 || M || y2)
    size_t ct3_temp_len = position_len * 2 + msg_len;
    uint8_t *ct3_temp = (uint8_t *)malloc(sizeof(uint8_t) * ct3_temp_len);
    uint8_t ct3[32];

    memcpy(ct3_temp, x2_y2, position_len);
    memcpy(ct3_temp + position_len, msg, msg_len);
    memcpy(ct3_temp + position_len + msg_len, x2_y2 + position_len, position_len);

    // print_chars_by_hex("ct3_temp", ct3_temp, ct3_temp_len);

    sm3_context ctx;

    sm3_init(&ctx);

    sm3_update(&ctx, ct3_temp, ct3_temp_len);

    sm3_finish(ct3, &ctx);

    // 4.输出结果
    // ct = C1 || C3 || C2
    memcpy(ct + ct_index, ct3, 32);
    ct_index = ct_index + 32;

    memcpy(ct + ct_index, ct2, msg_len);
    ct_index = ct_index + msg_len;

    // 5.释放空间
    free(ct3_temp);

    free(x2_y2);

    free(ct2);
}

/**
 * @brief SM2 解密函数
 *
 * @param plaintext 明文
 * @param plaintext_len 明文长度
 * @param sm2_param SM2参数
 * @param private_key 私钥
 * @param ct 密文
 * @param ct_len 密文长度
 */
void sm2_decrypt(uint8_t *plaintext, size_t plaintext_len, param sm2_param, mpz_t private_key, uint8_t *ct, size_t ct_len)
{
    // TODO 1.验证是否满足方程上的点，不满足退出
    // TODO 2.验证点 S=H[C1],若 S 是无穷远点 ,则 报错并退出;
    // 3.计算 db[C1]=private_key[C1]
    point k_g, private_key_c1;
    point_inits(k_g, private_key_c1, NULL);

    // 恢复k_g
    // 坐标的字节数
    size_t position_len = 24;
    // ct的索引
    size_t ct_index = 1;

    mpz_import(k_g->x, position_len / sizeof(mp_limb_t), 1, sizeof(mp_limb_t), 1, 0, ct + ct_index);
    ct_index = ct_index + position_len;
    // gmp_printf("kg=%ZX\n", k_g->x);

    mpz_import(k_g->y, position_len / sizeof(mp_limb_t), 1, sizeof(mp_limb_t), 1, 0, ct + ct_index);
    ct_index = ct_index + position_len;

    // 计算 db[C1]=private_key[C1]
    point_mul(private_key_c1, k_g, private_key, sm2_param->a, sm2_param->p);

    // 4.计算 t=KDF(x2||y2,klen)，若t全为0，报错退出
    size_t x2_y2_len = 2 * position_len;
    uint8_t *x2_y2 = (uint8_t *)malloc(sizeof(uint8_t) * x2_y2_len);

    mpz_to_byte_array(x2_y2, position_len, private_key_c1->x);
    mpz_to_byte_array(x2_y2 + position_len, position_len, private_key_c1->y);
    // print_chars_by_hex("x2_y2", x2_y2, x2_y2_len);

    kdf(plaintext, plaintext_len, x2_y2, x2_y2_len);
    // print_chars_by_hex("KDF plaintext", plaintext, plaintext_len);

    // 5.计算出 M'=C2 ^ t
    ct_index = ct_index + 32;
    for (size_t i = 0; i < plaintext_len; i++)
    {
        plaintext[i] = plaintext[i] ^ ct[i + ct_index];
    }

    // TODO 6.计算出 u = hash(x2 || M1 || y2)，若 u != C3 退出

    // 释放空间
    free(x2_y2);
}

/**
 * @brief 拼接信息
 *
 * @param out 输出字符串
 * @param out_len 输出字符串的长度
 * @param ent_len entlen_A转换而成的两个字节
 * @param id ID
 * @param id_len ID长度
 * @param sm2_param
 * @param base_point
 * @param public_key
 */
void concat_base_info(uint8_t *out, size_t out_len, size_t ent_len, uint8_t *id, size_t id_len, param sm2_param, point base_point, point public_key)
{
    // 用于将size_t转换为8个字节
    uint8_t size_t_tmp[8];
    // 记录当前out的索引
    size_t out_index = 0;
    // 坐标的字节数
    size_t position_len = 32;
    uint8_t position[position_len];

    // 1.拼接ent_len
    // 将ent_len转换8个字节，并取其低位的两个字节。注意这里是大端模式
    PUT_ULONG_BE(ent_len, size_t_tmp, 0);

    out[out_index++] = size_t_tmp[6];
    out[out_index++] = size_t_tmp[7];

    // 2.拼接id
    memcpy(out + out_index, id, id_len);
    out_index += id_len;

    // 3.拼接a，b
    mpz_to_byte_array(position, position_len, sm2_param->a);
    memcpy(out + out_index, position, position_len);
    out_index += position_len;

    mpz_to_byte_array(position, position_len, sm2_param->b);
    memcpy(out + out_index, position, position_len);
    out_index += position_len;

    // 4.拼接xG，yG
    mpz_to_byte_array(position, position_len, base_point->x);
    memcpy(out + out_index, position, position_len);
    out_index += position_len;

    mpz_to_byte_array(position, position_len, base_point->y);
    memcpy(out + out_index, position, position_len);
    out_index += position_len;

    // 5.拼接xA，yA
    mpz_to_byte_array(position, position_len, public_key->x);
    memcpy(out + out_index, position, position_len);
    out_index += position_len;

    mpz_to_byte_array(position, position_len, public_key->y);
    memcpy(out + out_index, position, position_len);
}

/**
 * @brief SM2签名函数
 *
 * @param sign 签名值
 * @param sign_len 签名值长度
 * @param sm2_param SM2参数
 * @param base_point G
 * @param private_key 私钥
 * @param base_info 基本信息。Z_A
 * @param base_info_len 基本信息的长度
 * @param msg 消息
 * @param msg_len 消息的长度
 */
void sm2_sign(uint8_t *sign, size_t sign_len, param sm2_param, point base_point, mpz_t private_key, uint8_t *base_info, size_t base_info_len, uint8_t *msg, size_t msg_len)
{
    // 公用变量
    mpz_t temp;
    mpz_init(temp);

    // 椭圆曲线上点坐标的字节数
    size_t position_len = 32;
    uint8_t position[position_len];
    // 字符串使用16进制
    size_t base = 16;

    // 1.拼接new_msg
    size_t new_msg_len = base_info_len + msg_len;
    uint8_t *new_msg = malloc(sizeof(uint8_t) * new_msg_len);

    memcpy(new_msg, base_info, base_info_len);
    memcpy(new_msg + base_info_len, msg, msg_len);

    // 2.计算哈希值e。e是new_msg的哈希值
    uint8_t new_msg_hash[33];
    hash(new_msg_hash, new_msg, new_msg_len);
    new_msg_hash[33] = '\0';

    mpz_t mpz_new_msg_hash;
    mpz_init(mpz_new_msg_hash);
    mpz_import(mpz_new_msg_hash, position_len / sizeof(mp_limb_t), 1, sizeof(mp_limb_t), 1, 0, new_msg_hash);

    // 3.使用随机数发生器生成随机数k，k∈[1,n-1]
    // 6CB28D99 385C175C 94F94E93 4817663F C176D925 DD72B727 260DBAAE 1FB2F96F
    mpz_t rand_k;
    mpz_init_set_str(rand_k, "6CB28D99385C175C94F94E934817663FC176D925DD72B727260DBAAE1FB2F96F", base);

    // 4.计算[k]G = [k] * base_point = (x1, y1)
    // 4.1.初始化多倍点 k[G]
    point k_g;
    point_init(k_g);

    // 4.2.计算 (x1, y1) = k[G] =rand_k[G]
    point_mul(k_g, base_point, rand_k, sm2_param->a, sm2_param->p);
    gmp_printf("k_g.x=%ZX\n", k_g->x);
    gmp_printf("k_g.y=%ZX\n", k_g->y);

    // 5.计算 r = (e + x1) mod n
    mpz_t r;
    mpz_init(r);

    mpz_add(r, mpz_new_msg_hash, k_g->x);
    mpz_mod(r, r, sm2_param->n);
    // gmp_printf("r=%ZX\n", r);

    // 6.计算 s = (k - r * d_A) * (1 + d_A)^(-1) mod n
    // 6.1.计算 (k - r * d_A) mod n
    mpz_t one_part;
    mpz_init(one_part);

    // 6.1.1.(r * d_A ) mod n = (r * private_key)  mod n
    mpz_t r_private_key;
    mpz_init(r_private_key);
    mpz_mul(r_private_key, r, private_key);
    mpz_mod(r_private_key, r_private_key, sm2_param->n);

    // 6.1.2.(k - r * d_A) mod n
    mpz_sub(one_part, rand_k, r_private_key);
    mpz_mod(one_part, one_part, sm2_param->n);

    // 6.2.计算 (1 + d_A)^-1 mod n
    mpz_t two_part;
    mpz_init(two_part);

    mpz_set_ui(temp, 1);
    mpz_add(temp, temp, private_key);
    mpz_mod(temp, temp, sm2_param->n);

    inverse_by_p(two_part, temp, sm2_param->n);
    // mpz_invert(two_part, temp, sm2_param->p);

    // 6.3.计算 s = (k - r * d_A) * (1 + d_A)^(-1) mod n
    mpz_t s;
    mpz_init(s);

    mpz_mul(s, one_part, two_part);
    mpz_mod(s, s, sm2_param->n);
    // gmp_printf("s=%ZX\n", s);

    // 7.将r、s转换为字符串输出
    // 7.1.将r转换为字符串，并放到sign中
    mpz_to_byte_array(position, position_len, r);
    memcpy(sign, position, position_len);

    // 7.2.将s转换为字符串，并放到sign中
    mpz_to_byte_array(position, position_len, s);
    memcpy(sign + position_len, position, position_len);

    // 8.收尾操作。例如，释放空间
    mpz_clears(temp, mpz_new_msg_hash, rand_k, NULL);

    point_clear(k_g);

    mpz_clears(r, one_part, r_private_key, two_part, s, NULL);
}

/**
 * @brief SM2签名函数
 *
 * @param sm2_param SM2参数
 * @param base_point 基点G
 * @param public_key 公钥
 * @param base_info 基本信息（接收到的信息）
 * @param base_info_len 基本信息长度
 * @param msg 消息（接收到的消息）
 * @param msg_len 消息的长度
 * @param sign 签名值
 * @param sign_len 签名值长度
 * @return int 返回1表示验签成功，返回其他值表示验签失败
 */
int sm2_verify_sign(param sm2_param, point base_point, point public_key, uint8_t *base_info, size_t base_info_len, uint8_t *msg, size_t msg_len, uint8_t *sign, size_t sign_len)
{
    // 公用变量
    mpz_t temp_mpz;
    mpz_init(temp_mpz);
    point temp_point;
    point_init(temp_point);

    // 字符串使用16进制
    size_t base = 16;
    // 椭圆曲线上点坐标的字节数
    size_t position_len = 32;
    uint8_t position[position_len];
    uint8_t temp_chars[position_len + 1];
    temp_chars[position_len + 1] = '\0';

    // 返回值
    int ret = 0;

    // TODO 1.检验 r'∈[1,n-1]是否成立，若不成立表示验证不通过
    // TODO 2.检验 s'∈[1,n-1]是否成立，若不成立表示验证不通过

    // 3.计算new_msg
    size_t new_msg_len = base_info_len + msg_len;
    uint8_t *new_msg = malloc(sizeof(uint8_t) * new_msg_len);
    memcpy(new_msg, base_info, base_info_len);
    memcpy(new_msg + base_info_len, msg, msg_len);

    // 4.计算哈希值e。e是new_msg的哈希值
    uint8_t new_msg_hash[32];
    hash(new_msg_hash, new_msg, new_msg_len);

    // 5.将r'，s'转换为整数，计算 t =(r' + s') mod n，若 t = 0，则验证不通过
    // 5.1.将r'，s'转换整数
    mpz_t r, s, t;
    mpz_inits(r, s, t, NULL);

    memcpy(temp_chars, sign, position_len);
    // mpz_init_set_str(r, temp_chars, base);
    mpz_import(r, position_len / sizeof(mp_limb_t), 1, sizeof(mp_limb_t), 1, 0, temp_chars);

    memcpy(temp_chars, sign + position_len, position_len);
    // mpz_init_set_str(s, temp_chars, base);
    mpz_import(s, position_len / sizeof(mp_limb_t), 1, sizeof(mp_limb_t), 1, 0, temp_chars);

    // 5.2.计算 t =(r' + s') mod n
    mpz_add(t, r, s);
    mpz_mod(t, t, sm2_param->n);
    gmp_printf("t=%ZX\n", t);

    // 5.3.若 t = 0，则验证不通过

    // 6.计算(x1', y1') = [s']G + [t]PA
    // 6.1.计算[s']G = [s'] * base_point
    point s_g;
    point_init(s_g);

    point_mul(s_g, base_point, s, sm2_param->a, sm2_param->p);
    gmp_printf("s_g.x=%ZX\n", s_g->x);

    // 6.2.计算[t]P_A =[t] * public_key
    point t_public_key;
    point_init(t_public_key);
    point_mul(t_public_key, public_key, t, sm2_param->a, sm2_param->p);

    gmp_printf("t_public_key.x=%ZX\n", t_public_key->x);

    // 6.3.计算(x1', y1') = [s']G + [t]PA
    point_add(temp_point, s_g, t_public_key, sm2_param->a, sm2_param->p);

    mpz_t x1;
    mpz_init(x1);
    // mpz_add(x1, s_g->x, t_public_key->x);
    mpz_set(x1, temp_point->x);
    gmp_printf("x1=%ZX\n", x1);

    // 7.计算 R = (e' + x1') mod n，检查R = r'是否成立，若成立则验证成功
    // 7.1.将e'转为整数
    mpz_t mpz_new_msg_hash;
    mpz_init(mpz_new_msg_hash);
    memcpy(temp_chars, new_msg_hash, position_len);
    // mpz_init_set_str(mpz_new_msg_hash, temp_chars, base);
    mpz_import(mpz_new_msg_hash, position_len / sizeof(mp_limb_t), 1, sizeof(mp_limb_t), 1, 0, temp_chars);

    // 7.2.计算 R = (e' + x1') mod n
    mpz_t new_r;
    mpz_init(new_r);
    mpz_add(new_r, mpz_new_msg_hash, x1);
    mpz_mod(new_r, new_r, sm2_param->n);

    // 7.3.检查R = r'是否成立，若成立则验证成功
    if (mpz_cmp(new_r, r) == 0)
    {
        ret = 1;
    }

    // 8.收尾操作。例如，释放空间
    mpz_clears(temp_mpz, r, s, t, NULL);

    point_clears(temp_point, t_public_key, NULL);

    mpz_clears(x1, mpz_new_msg_hash, new_r, NULL);

    return ret;
}

/**
 * @brief 哈希消息
 *
 * @param out 输出 SM3输出是固定的32个字节
 * @param msg 消息
 * @param msg_len 消息长度
 */
void hash(uint8_t *out, uint8_t *msg, size_t msg_len)
{
    sm3_context ctx;

    sm3_init(&ctx);

    sm3_update(&ctx, msg, msg_len);

    sm3_finish(out, &ctx);
}
/**
 * @brief mpz转byte数组
 *
 * @param byte_array 输出
 * @param size 数组的长度，包括填充
 * @param value mpz的值
 */
void mpz_to_byte_array(unsigned char *byte_array, size_t size, mpz_t value)
{
    size_t n, limbs;

    // 获取所需的字节数量
    n = mpz_sizeinbase(value, 256);

    // 获取 limb 数量
    limbs = (n + (sizeof(mp_limb_t) * 8) - 1) / (sizeof(mp_limb_t) * 8);

    // 导出到字节数组
    // 输出数组、转为256进制之后的长度，1表示不逆序，sizeof(mp_limb_t) 一个limb的大小，1小端，钉子位空间够默认为0，输出结果value

    // 计算填充位置
    size_t padding_len = size - n;
    if (padding_len > 0)
    {
        mpz_export(byte_array + padding_len, &n, 1, sizeof(mp_limb_t), 1, 0, value);
    }
    else
    {
        mpz_export(byte_array, &n, 1, sizeof(mp_limb_t), 1, 0, value);
    }
}

/**
 * 将mpz转换为固定位数的字符串.
 *
 * \param num mpz格式的数
 * \param length 输出的字符串的长度
 * \param base 基数
 * \return
 */
void mpz_to_fixed_length_str(char *padded_str, size_t padded_str_len, mpz_t num, int base)
{
    // 先将 mpz 的数据转换为 str
    char *raw_str = malloc(sizeof(char) * padded_str_len);

    mpz_get_str(raw_str, base, num);

    int raw_len = strlen(raw_str);

    // 计算需要填充的零的数量
    // -1 是因为有一个长度给结束符号的
    int padding_len = padded_str_len - 1 - raw_len;
    if (padding_len <= 0)
    {
        // 实际长度大于等于所需长度，无需填充
        strcpy(padded_str, raw_str);

        // 释放原始字符串占用的内存
        free(raw_str);

        return;
    }

    // 填充零
    memset(padded_str, '0', padding_len);

    // 复制原始数字字符串到填充后的字符串中
    strcpy(padded_str + padding_len, raw_str);

    // 释放原始字符串占用的内存
    free(raw_str);
}